import Yices hiding (and, or, not)
import qualified Yices as Y
import Prelude hiding (sum)

-- Goods
type Truck = Nat
data Goods = Cheese | Beer | Wine | Drinks | Chips deriving (Bounded, Enum, Show)
instance Named Goods where sname _ = "goods"
instance Refl Goods where refl = stdRefl
instance Scalar Goods

-- Parameters
minBeerCount   = 19
maxTruckCount  = 8
maxTruckWeight = 8500
truckCount     = 6

trucks = [1..truckCount] :: [Truck]
goods = [minBound..maxBound] :: [Goods]

-- Weight of a pallet of different goods
weight :: Goods -> Nat
weight Cheese = 700
weight Beer   = 1300
weight Wine   = 1000
weight Drinks = 1500
weight Chips  = 100

-- Variables
load' :: Ident (Goods -> Truck -> Nat)
(load',load) = ident2 "load"

-- Predicates
sum :: [Term Nat] -> Term Nat
sum = foldl1 (\x y -> idt plus' ! x ! y)

totalPallets :: Goods -> Term Nat
totalPallets g = sum [load g t | t <- trucks]

totalLoad :: Truck -> Term Nat
totalLoad t = sum [load g t * weight g | g <- goods]
  where
    x * y = mul' ! x ! nat y

totalCount :: Truck -> Term Nat
totalCount t = sum [load g t | g <- goods]

script =
    -- Define the used types and variables
    [ deftype (undefined :: Goods)
    , define load'
    -- Now just list the requirements
    , assert $ totalPallets Cheese == 4
    , assert $ totalPallets Wine   == 8
    , assert $ totalPallets Drinks == 10
    , assert $ totalPallets Chips  == 5
    , assert $ totalPallets Beer   >= minBeerCount
    -- Make sure we do not overload the trucks
    , assert $ conj [totalLoad t <= maxTruckWeight | t <- trucks]
    , assert $ conj [totalCount t <= maxTruckCount | t <- trucks]
    ]
  where
    (<=) x y = le' ! x ! nat y
    (==) x y = eq' ! x ! nat y
    (>=) x y = ge' ! x ! nat y

main = printScript script
